package me.itsoo.artemis.framework.repository.datasource;

import lombok.Getter;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import me.itsoo.artemis.framework.core.util.Assert;
import me.itsoo.artemis.framework.core.util.IoUtils;
import me.itsoo.artemis.framework.core.util.StringUtils;
import me.itsoo.artemis.framework.repository.config.Configuration;
import me.itsoo.artemis.framework.repository.exception.CreatedConnectionException;
import me.itsoo.artemis.framework.repository.exception.UncategorizedSqlException;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.Objects;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * DataSourceManager
 *
 * @author zxy
 * @version 0.0.1
 * @since 2023/5/18
 */
@Slf4j
public class DataSourceManager implements ConnectionPool {

    @Getter
    private final Configuration config;

    private final SqlConnectionProxyFactory connectionFactory;

    private final LinkedBlockingQueue<SqlConnection> activePool;

    private final ScheduledExecutorService sec;

    private final AtomicBoolean secSwitch;

    private final AtomicInteger createdCounter;

    /**
     * 核心连接数
     */
    private final int coreSize;

    /**
     * 最大连接数
     */
    private final int maxSize;

    /**
     * 最大等待连接时间（指获取连接动作）
     */
    private final int maxWait;

    /**
     * 连接保活时间（秒）
     */
    private final int keepAlive;

    /**
     * 连接保活 SQL
     */
    private final String validationQuery;

    public DataSourceManager(Configuration config) {
        Assert.notNull(config, "'config' cannot be null");

        this.config = config;
        this.connectionFactory = new SqlConnectionProxyFactory(this);
        this.activePool = new LinkedBlockingQueue<>(config.getCoreSize());
        this.sec = new ScheduledThreadPoolExecutor(1, new DefaultScheduledThreadFactory());
        this.secSwitch = new AtomicBoolean(false);
        this.createdCounter = new AtomicInteger(0);

        this.coreSize = config.getCoreSize();
        this.maxSize = config.getMaxSize();
        this.maxWait = config.getMaxWait();
        this.keepAlive = config.getKeepAlive();
        this.validationQuery = config.getValidationQuery();

        init();
    }

    @SneakyThrows(ClassNotFoundException.class)
    private void init() {
        Class.forName(config.getDriverClassName());

        for (int i = 0; i < coreSize; i++) {
            activePool.add(getConnection0());
            createdCounter.incrementAndGet();
        }

        listen();
    }

    @Override
    public SqlConnection getConnection(boolean autoCommit) {
        final SqlConnection result = getConnection();

        try {
            result.setAutoCommit(autoCommit);
        } catch (SQLException e) {
            throw new UncategorizedSqlException("No supported auto-commit");
        }

        if (log.isDebugEnabled()) {
            log.debug("<CONNECT> [{}] curr_count: {}, curr_free: {}, total: {}",
                    Thread.currentThread().getId(), createdCounter.get(), activePool.size(), maxSize);
        }

        return result;
    }

    @Override
    public SqlConnection getConnection() {
        SqlConnection result = activePool.poll();
        if (Objects.nonNull(result)) {
            return result;
        }

        if (createdCounter.incrementAndGet() <= maxSize) {
            result = getConnection0();
            return result;
        } else {
            // 未获取到连接需要计数器 -1
            createdCounter.decrementAndGet();
        }

        try {
            result = activePool.poll(maxWait, TimeUnit.MILLISECONDS);
            if (Objects.nonNull(result)) {
                return result;
            }
            throw new InterruptedException("fire");
        } catch (InterruptedException e) { // NOSONAR
            throw new CreatedConnectionException(coreSize, maxSize);
        }
    }

    @Override
    public SqlConnection getConnection(String username, String password) {
        try {
            final Connection conn = DriverManager.getConnection(config.getUrl(), username, password);
            return connectionFactory.getSqlConnection(conn);
        } catch (SQLException e) {
            final String errorMsg = StringUtils.format("Fail: creating conn, url: {}, user: {}, password: {}",
                    config.getUrl(), username, password);
            throw new CreatedConnectionException(errorMsg);
        }
    }

    private SqlConnection getConnection0() {
        return getConnection(config.getUsername(), config.getPassword());
    }

    @Override
    public void releaseConnection(SqlConnection conn) {
        if (createdCounter.decrementAndGet() >= coreSize) {
            IoUtils.close(conn.getSourceConnection());
        } else {
            activePool.add(conn);
        }

        if (log.isDebugEnabled()) {
            log.debug("<RELEASE> [{}] curr_count: {}, curr_free: {}, total: {}",
                    Thread.currentThread().getId(), createdCounter.get(), activePool.size(), maxSize);
        }
    }

    @Override
    public boolean isActiveConnection(SqlConnection conn) {
        return Objects.nonNull(conn) && activePool.contains(conn);
    }

    /**
     * 保活监听
     */
    private void listen() {
        while (!secSwitch.get()) {
            if (secSwitch.compareAndSet(false, true)) {
                sec.scheduleWithFixedDelay(this::probingAll, keepAlive, keepAlive, TimeUnit.SECONDS);
            }
        }
    }

    /**
     * 连接保活
     */
    @SuppressWarnings("squid:S3776")
    private void probingAll() {
        final long currTimestamp = System.currentTimeMillis();
        final long keepAliveMillis = keepAlive * 1000L;

        for (int i = 0, checkCount = activePool.size(); i < checkCount; i++) {
            final SqlConnection conn = activePool.poll();
            if (Objects.nonNull(conn) && !conn.isClosed()) {
                final long diff = (currTimestamp - conn.getLastUsedTimestamp());
                if (diff >= keepAliveMillis) {
                    try {
                        if (StringUtils.isBlank(validationQuery)) {
                            probingPing(conn, currTimestamp);
                        } else {
                            probingSql(conn, currTimestamp);
                        }
                    } catch (SQLException | RuntimeException ignore) {
                        // Ignore connection timeout
                    }
                }
            }
        }
    }

    /**
     * PING 方式探活
     *
     * @param conn          SqlConnection
     * @param currTimestamp long
     */
    private void probingPing(SqlConnection conn, long currTimestamp) {
        if (conn.isValid()) {
            conn.setLastUsedTimestamp(currTimestamp);
            activePool.add(conn);
        }
    }

    /**
     * SQL 方式探活
     *
     * @param conn          SqlConnection
     * @param currTimestamp long
     * @throws SQLException SQLException
     */
    private void probingSql(SqlConnection conn, long currTimestamp) throws SQLException {
        final Connection sc = conn.getSourceConnection();
        try (Statement stmt = sc.createStatement()) {
            try (ResultSet rs = stmt.executeQuery(validationQuery)) {
                if (rs.next()) {
                    conn.setLastUsedTimestamp(currTimestamp);
                    activePool.add(conn);
                }
            }
        }
    }

    /**
     * DefaultScheduledThreadFactory
     */
    private static class DefaultScheduledThreadFactory implements ThreadFactory {

        private final AtomicInteger threadNumber = new AtomicInteger(1);

        @Override
        public Thread newThread(Runnable r) {
            final Thread result = new Thread(r);
            if (result.isDaemon()) {
                result.setDaemon(false);
            }
            if (result.getPriority() != Thread.NORM_PRIORITY) {
                result.setPriority(Thread.NORM_PRIORITY);
            }

            result.setName("datasource-thread-" + threadNumber.getAndIncrement());
            return result;
        }
    }
}
